import numpy
import pylab
import cPickle
import general_neuro.fast_thresh_detect as fast_thresh_detect
import general_neuro.Iion as Iion
from genutils.smoothing import adc_interp

print "loading in cPickle"
data_set_list = cPickle.load(open('intracellular_current_injection_fitting_data_FINAL.cPickle', 'r'))
print 'data loaded..'

def plot_data_set(count, data):
    # also plots spike times as axvline
    pylab.figure(0)
    if count:
        pylab.clf()
    # plot the data
    num_data = len(data.keys())
    num_rows = int(numpy.sqrt(num_data))
    num_cols = int(numpy.sqrt(num_data))+1
    if num_rows * num_cols < num_data:
        num_rows += 1
    plot_count = 0

    data_list   = []
    inject_list = []
    title_list  = []
    isi_list    = []
    times_list  = []
    for data_title, data in data.items():
        data_list.append(data['data'])
        inject_list.append(data['inject'])
        times_list.append(data['abs_spike_index'])
        title_list.append(data_title)
        isi_list.append(data['isi'])
        Iion_list.append(data["Iion"])
        Iinj_list.append(data['Iinj'])
    indexes = numpy.argsort(inject_list)

    for i in indexes:
        plot_count += 1
        pylab.figure(0)
        pylab.subplot(num_rows,num_cols, plot_count)
        pylab.plot(data_list[i])
        pylab.plot(data_list[i])
        pylab.plot(Iion_list[i]*10) # scaling up so we can see it
        pylab.plot(Iinj_list[i]*10)
        pylab.xticks('')
        pylab.yticks('')
        pylab.xlabel('%1.2f nA' % inject_list[i])
        pylab.ylabel(title_list[i][-10:])
        for spike_index in times_list[i]:
            pylab.axvline(spike_index, color='red')
    pylab.suptitle(title_list[i].split('--')[0])
    pylab.figure(1)
    pylab.clf()
    for isi_set in isi_list:
        pylab.plot(isi_set, '-o')

count = 0
for data_set in data_set_list:
    # say what set of data we're finding spike times in.
    print "finding spikes for %s" % data_set.keys()[0].split('--')[0]
    C_list = []
    for name, data in data_set.items():
        if data['inject'] > 0.0:
            data['abs_spike_index'], junk = \
               fast_thresh_detect.fast_thresh_detect( data['data'], threshold=-10.0)
            # get rid of the last threshold crossing since it is an artifact.
            data['abs_spike_index'] = numpy.array( data['abs_spike_index'][:-1] )
            print "Found %d spikes in %s." % (len(data['abs_spike_index']), name)
            print data['abs_spike_index']
            # find the inter-spike-intervals assuming dt = 0.1 ms
            si_list = data['abs_spike_index']
            isi = []
            for i in xrange(len(si_list)-1):
                isi.append((si_list[i+1]-si_list[i])*0.1)
            data['isi'] = numpy.array(isi)
            # find spike times relative to current pulse onset
            pulse_onset, junk = fast_thresh_detect.fast_thresh_detect( data['stim2'], threshold=0.1)
            rst_list = [(si_list[i]-pulse_onset[0])*0.1 for i in xrange(len(si_list))]
            data['relative_spike_times'] = numpy.array(rst_list)
        else:
            data['abs_spike_index'] = []
            data['relative_spike_times'] = []
            data['isi'] = []
        if "C" in data.keys():
            C_list.append(data['C'])
    if C_list:
        for name, data in data_set.items():
            data['Cavg'] = numpy.average(C_list)
            # correctly scale the injected current.
            data['Iinj'] = abs(data['stim2'])/max(abs(data['stim2']))*data['inject']
            del(data['stim2'])
    count += 1

print "saving cPickle"
cPickle.dump(data_set_list, open('intracellular_current_injection_fitting_data_FINAL_.cPickle', 'w'), protocol=-1)
print 'data saved..'
